import numpy as np
from BaselineRemoval import BaselineRemoval
from matplotlib import pyplot as plt
from input import input_data

def polynomial_baseline_correction(data:np.ndarray, degree=3):
    '''
    对一维数据进行多项式基线校正
    :param data: 需要进行校正的一维数据数组。np.ndarray nxm
    :param degree: 用于拟合的多项式的次数。默认为3。int
    :return: 去除基线后的校正数据。 np.ndarray
    '''
    corrected_data = np.array([baseline_correction(row, degree) for row in data])
    return corrected_data


def baseline_correction(row, degree=3):
    '''
    接受一个一维数组并返回校正后的结果
    :param row:
    :param degree:
    :return:
    '''
    baseObj = BaselineRemoval(row.tolist())
    Modpoly_output = baseObj.ModPoly(degree)
    # 另外两种内置的方法
    # Imodpoly_output = baseObj.IModPoly(degree)
    # Zhangfit_output = baseObj.ZhangFit()
    return np.array(Modpoly_output)

def polynomial_baseline_correction2(data:np.ndarray, degree=3):
    '''
    手动实现：对一维数据进行多项式基线校正
    :param data: 需要进行校正的一维数据数组。np.ndarray
    :param degree: 用于拟合的多项式的次数。默认为3。int
    :return: 去除基线后的校正数据。 np.ndarray
    '''
    # 创建数据的索引数组
    x = np.arange(len(data))
    # 使用多项式拟合数据
    p = np.polyfit(x, data, degree)  # p 为多项式的系数
    # 计算多项式在各个数据点上的值
    y_poly = np.polyval(p, x)
    # 从原始数据中减去多项式值，得到校正后的数据
    corrected_data = data - y_poly
    return corrected_data

if __name__ == "__main__":
    data,_ = input_data()
    n = data.shape[1]
    x = np.linspace(1, n, n)
    # 创建一些带基线的数据
    #x = np.linspace(0, 10, 100)
    #y = 3 * x ** 2 + 2 * x + 10 + np.random.normal(0, 5, 100)  # 模拟的基线是3x^2 + 2x + 10，加上一些噪声
    # 进行基线校正
    corrected_y = polynomial_baseline_correction(data, degree=2)
    # 绘制原始数据和校正后的数据
    plt.figure(figsize=(10, 5))
    plt.plot(x, data[0], label='Original Data', color='blue')
    plt.plot(x, corrected_y[0], label='Corrected Data', color='red')
    plt.legend()
    plt.xlabel('X')
    plt.ylabel('Y')
    plt.title('Polynomial Baseline Correction')
    plt.show()